import copy
import pickle
import imageai
import os
import math
from concurrent.futures import ProcessPoolExecutor
from imageai.Detection import VideoObjectDetection
#from numba import jit

execution_path = os.getcwd()

readpath = r'G:\software\work\finalhomework\video'
#writepath = r'G:\software\work\finalhomework\frame'
color_index = {'bus': 'red', 'handbag': 'steelblue', 'giraffe': 'orange', 'spoon': 'gray', 'cup': 'yellow', 'chair': 'green', 
            'elephant': 'pink', 'truck': 'indigo', 'motorcycle': 'azure', 'refrigerator': 'gold', 'keyboard': 'violet', 'cow': 'magenta', 
            'mouse': 'crimson', 'sports ball': 'raspberry', 'horse': 'maroon', 'cat': 'orchid', 'boat': 'slateblue', 'hot dog': 'navy', 
            'apple': 'cobalt', 'parking meter': 'aliceblue', 'sandwich': 'skyblue', 'skis': 'deepskyblue', 'microwave': 'gold', 
            'knife': 'cadetblue', 'baseball bat': 'cyan', 'oven': 'lightcyan', 'carrot': 'coldgrey', 'scissors': 'seagreen', 'sheep': 'deepgreen', 
            'toothbrush': 'cobaltgreen', 'fire hydrant': 'limegreen', 'remote': 'forestgreen', 'bicycle': 'olivedrab', 'toilet': 'ivory', 
            'tv': 'khaki', 'skateboard': 'palegoldenrod', 'train': 'cornsilk', 'zebra': 'wheat', 'tie': 'burlywood', 'orange': 'melon', 
            'bird': 'bisque','dining table': 'chocolate', 'hair drier': 'sandybrown', 'cell phone': 'sienna', 'sink': 'coral', 'bench': 'salmon', 
            'bottle': 'brown', 'car': 'silver', 'bowl': 'maroon', 'tennis racket': 'palevilotered', 'airplane': 'lavenderblush', 'pizza': 'hotpink', 
            'umbrella': 'deeppink', 'bear': 'plum', 'fork': 'purple', 'laptop': 'indigo', 'vase': 'mediumpurple', 'baseball glove': 'slateblue', 
            'traffic light': 'mediumblue', 'bed': 'navy', 'broccoli': 'royalblue', 'backpack': 'slategray', 'snowboard': 'skyblue', 'kite': 'cadetblue', 
            'teddy bear': 'orange', 'clock': 'lightcyan', 'wine glass': 'teal', 'frisbee': 'aquamarine', 'donut': 'mincream', 'suitcase': 'seagreen', 
            'dog': 'springgreen', 'banana': 'emeraldgreen', 'person': 'honeydew', 'surfboard': 'palegreen', 'cake': 'sapgreen', 'book': 'lawngreen', 
            'potted plant': 'greenyellow', 'toaster': 'ivory', 'stop sign': 'beige', 'couch': 'khaki'}

def inititems():
    '''
    初始化
    '''
    raw_items = {}
    for item in color_index:
        raw_items[item] = 0
    return raw_items

items = inititems()

def forframe(frame_number, output_array, output_count, returned_frame):
    '''
    对每一帧，提取物体出现的次数
    '''
    if(output_count == {}):
        return
    global items
    for eachItem in output_count:
        items[eachItem] += output_count[eachItem]

def forFull(output_arrays, count_arrays, average_output_count):
    print("Array for the outputs of each frame ", output_arrays)
    print("Array for output count for unique objects in each frame : ", count_arrays)
    print("Output average count for unique objects in the entire video: ", average_output_count)
    print("------------END OF THE VIDEO --------------")

def detect(*argv):
    '''
    检测一个视频，返回检测出的实体字典
    '''
    print(argv)
    root = argv[0]
    file = argv[1]
    i = argv[2]
    global items
    items=inititems()
    print('-------------------------{}--------------------------'.format(i))
    # items=raw_item
    detector = VideoObjectDetection()
    detector.setModelTypeAsTinyYOLOv3()
    detector.setModelPath(os.path.join(execution_path , r"video_detection\yolo-tiny.h5"))
    detector.loadModel(detection_speed='flash')
    video = detector.detectObjectsFromVideo(input_file_path=os.path.join(root, file),  #注意不同电脑的路径要修改一下
                                            save_detected_video=False,
                                            #video_complete_function=forFull,
                                            minimum_percentage_probability=30, per_frame_function=forframe,
                                            return_detected_frame=True)  
    return items

if __name__ == '__main__':
    pass
    